#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2022/11/8 22:55
# @USER    : Shengji He
# @File    : visual.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
from PIL import Image
import os, glob
import torch
from torchvision.utils import save_image, make_grid


def display():
    folder = './run/samples/baseline_mnist'
    glob_research = os.path.join(folder, '*.pth')
    files = glob.glob(glob_research)
    grid_size = 10
    imgs = []

    for i, file in enumerate(files):
        sample = torch.load(file,)
        image_grid = make_grid(sample, nrow=grid_size)
        if i % 10 == 0:
            im = Image.fromarray(
                image_grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy())
            imgs.append(im)

        save_image(image_grid, os.path.join(folder, 'image_{}.png'.format(i)))

    imgs[0].save(os.path.join(folder, "movie.gif"), save_all=True, append_images=imgs[1:], duration=1, loop=0)
    pass


if __name__ == '__main__':
    display()
    print('done')
